import zipfile, os
import subprocess
import shutil


def update(update_file_path, save_root_dir, backup_dir):
    """ 更新程序
    args:
        update_file_path: zip升级包路径
        save_root_dir: 需要覆盖的软件根目录
        back_up_dir: 备份路径, 为None则不进行备份

    zip升级包结构:
    update_file.zip:
        del_file.txt # 记录了哪些文件需要被删除, 该文件可以不存在, 不存在时表示没有文件被删除
        xxx          # 项目下文件, 更新包下的根目录与项目根目录等价
        
    备份路径格式为:
    back_up_dir:
        add_file.txt    保留了本次更新中新增的文件
        files           这是一个文件夹, 在更新后被替换/删除的文件都会保留在这里
    
    """
    # 清除上次更新的备份
    shutil.rmtree(backup_dir,ignore_errors=True)
    os.makedirs(backup_dir)
    # add_file.txt
    add_file = open("{}/add_file.txt".format(backup_dir),"a")

    # unzip读取压缩包内容
    zip_file = zipfile.ZipFile(update_file_path)
    path_list = zip_file.namelist() # 获取压缩包内的文件遍历
    
    for path_ in path_list:
        if path_[-1] == os.path.sep:
            continue
        if os.path.basename(path_) == "del_file.txt":
            continue
        # 如果原文件存在, 则进行备份
        org_file_path = os.path.join(save_root_dir,path_)
        if os.path.exists(org_file_path):
            org_file_path = os.path.join(save_root_dir,path_)
            backup_file_path = os.path.join(backup_dir,"files",path_)
            if not os.path.exists(os.path.dirname(backup_file_path)):
                os.makedirs(os.path.dirname(backup_file_path))
            shutil.move(org_file_path, backup_file_path)
        else: # 如果文件不存在, 则说明是新增文件, 记录在change_file中
            add_file.write("{}\n".format(path_))
        # 将新文件解压
        zip_file.extract(path_,save_root_dir)


# 回滚操作
def roll_back(save_root_dir, backup_dir):
    add_files_path = backup_dir+"/add_file.txt"
    # 如果这个新增文件记录不存在, 说明没有可回滚的记录, 无法回滚
    if not os.path.exists(add_files_path):
        return False
    # 将新增的文件删除
    with open(add_files_path,"r") as f:
        for readline in f.readlines():
            readline = readline.strip()
            add_file_path = save_root_dir+"/"+readline
            os.remove(add_file_path)
    # 将备份的文件覆盖回去
    dirs = [x for x in os.walk(backup_dir+"/files")]
    for dir_ in dirs:
        for file_name in dir_[2]:
            file_path = dir_[0]+"/"+file_name
            
            target_path = file_path.split(backup_dir+"/files")[-1] # 这样切分出来的路径以/开头
            shutil.move(file_path,save_root_dir+target_path)

    # 最后清空备份文件夹
    shutil.rmtree(backup_dir,ignore_errors=True)
    os.makedirs(backup_dir)
    return True

if __name__ == "__main__":
    import os

    a = os.popen("ps -ef |grep simple_server.py")
    a = a.readlines()

    # update("data/update/upgrade.zip","./","data/update/backup")
    roll_back("./","data/update/backup")